{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eae90798fac43d38",
   "metadata": {},
   "source": [
    "# 手写推理大模型\n",
    "作者：IT周瑜\n",
    "\n",
    "微信ID：it_zhouyu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "965c51781e36b009",
   "metadata": {},
   "source": [
    "所谓推理大模型，就是在回答问题的时候有一个Think的过程，而之所以有这个过程，是因为在提供的训练数据中提供了Think过程的训练数据。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "613ca955598753ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:25:55.426650Z",
     "start_time": "2025-07-23T01:25:55.420778Z"
    }
   },
   "outputs": [],
   "source": [
    "zhouyu_chat_template = \"\"\"<|zhouyu_start|>user\n",
    "{question}<|zhouyu_end|>\n",
    "<|zhouyu_start|>assistant\n",
    "<think>\n",
    "{think}\n",
    "</think>\n",
    "{answer}\n",
    "<|zhouyu_end|>\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c80e7165d165de0b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:25:55.488247Z",
     "start_time": "2025-07-23T01:25:55.486131Z"
    }
   },
   "outputs": [],
   "source": [
    "data = [\n",
    "    {\n",
    "        \"question\": \"你叫什么名字？\",\n",
    "        \"think\": \"我得思考一下，嗯....，我知道了\",\n",
    "        \"answer\": \"我叫周瑜\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d5230c4ecd75319f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:25:55.500577Z",
     "start_time": "2025-07-23T01:25:55.496837Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|zhouyu_start|>user\\n你叫什么名字？<|zhouyu_end|>\\n<|zhouyu_start|>assistant\\n<think>\\n我得思考一下，嗯....，我知道了\\n</think>\\n我叫周瑜\\n<|zhouyu_end|>'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zhouyu_chat_template.format(question=data[0][\"question\"], think=data[0][\"think\"], answer=data[0][\"answer\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "11a007aab6708dc2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:01.047851Z",
     "start_time": "2025-07-23T01:25:55.528093Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dadudu/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/openai-community/gpt2\n",
      "Downloading Model from https://www.modelscope.cn to directory: /Users/dadudu/.cache/modelscope/hub/models/openai-community/gpt2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Embedding(50259, 768)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from modelscope import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "model_name = \"openai-community/gpt2\"\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "tokenizer.add_special_tokens({'bos_token': '<|zhouyu_start|>'})\n",
    "tokenizer.add_special_tokens({'eos_token': '<|zhouyu_end|>'})\n",
    "tokenizer.add_special_tokens({'pad_token': '<|endoftext|>'})\n",
    "\n",
    "model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "880f71b43ce0ba78",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:01.066442Z",
     "start_time": "2025-07-23T01:26:01.062574Z"
    }
   },
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "# 微信id：it_zhouyu\n",
    "class ZhouyuDataset(Dataset):\n",
    "    def __init__(self, data, max_length=256):\n",
    "        self.encodings = []\n",
    "        for qa in data:\n",
    "            text = zhouyu_chat_template.format(question=qa[\"question\"], think=qa[\"think\"], answer=qa[\"answer\"])\n",
    "            encoded = tokenizer(\n",
    "                text,\n",
    "                max_length=max_length,\n",
    "                padding='max_length',\n",
    "                truncation=True,\n",
    "                return_tensors='pt'\n",
    "            )\n",
    "            input_ids = encoded['input_ids'].squeeze()\n",
    "            self.encodings.append(input_ids)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.encodings)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.encodings[idx]\n",
    "\n",
    "\n",
    "dataset = ZhouyuDataset(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2e47fde54b6e976e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:01.076615Z",
     "start_time": "2025-07-23T01:26:01.072666Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([50257,  7220,   198, 19526,   254, 20998,   104, 20015,   222, 20046,\n",
       "          230, 28938,   235, 27764,   245,   171,   120,   253, 50258,   198,\n",
       "        50257,   562, 10167,   198,    27, 14925,    29,   198, 22755,   239,\n",
       "        36181,   245, 45250,   251, 32003,   225, 31660, 10310,   233,   171,\n",
       "          120,   234,   161,   245,   107,  1106,   171,   120,   234, 22755,\n",
       "          239,   163,   253,    98, 34402,   241, 12859,   228,   198,  3556,\n",
       "        14925,    29,   198, 22755,   239, 20998,   104, 37772,   101,   163,\n",
       "          239,   250,   198, 50258, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256, 50256,\n",
       "        50256, 50256, 50256, 50256, 50256, 50256])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "df53a9dcaa1535fe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:01.096655Z",
     "start_time": "2025-07-23T01:26:01.093920Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<|zhouyu_start|>user\\n你叫什么名字？<|zhouyu_end|>\\n<|zhouyu_start|>assistant\\n<think>\\n我得思考一下，嗯....，我知道了\\n</think>\\n我叫周瑜\\n<|zhouyu_end|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.decode(dataset[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b186d07e13acb6f4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:01.118433Z",
     "start_time": "2025-07-23T01:26:01.109145Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "\n",
    "# 创建数据收集器\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer,\n",
    "    mlm=False  # 使用CLM（因果语言模型）\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a79834f145c5b86f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:45.431128Z",
     "start_time": "2025-07-23T01:26:01.132077Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/dadudu/miniconda3/envs/mini-gpt/lib/python3.10/site-packages/torch/utils/data/dataloader.py:683: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used.\n",
      "  warnings.warn(warn_msg)\n",
      "`loss_type=None` was set in the config but it is unrecognised.Using the default loss: `ForCausalLMLoss`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='200' max='200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [200/200 00:42, Epoch 200/200]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>2.980600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>1.229300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>0.681700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>0.527200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>0.472900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>0.448100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>0.394500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>0.371700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>0.345400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.307100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>110</td>\n",
       "      <td>0.224900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>0.222100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>130</td>\n",
       "      <td>0.198600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>0.192100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>0.186600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>0.152200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>170</td>\n",
       "      <td>0.197300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>0.154000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>190</td>\n",
       "      <td>0.156000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.145500</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "# 训练配置\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./zhouyu_think_model\",\n",
    "    per_device_train_batch_size=1,\n",
    "    num_train_epochs=200,\n",
    "    # eval_strategy=\"epoch\",\n",
    "    # save_strategy=\"epoch\",\n",
    "    logging_steps=10\n",
    ")\n",
    "\n",
    "# 创建Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset,\n",
    "    # eval_dataset=tokenized_datasets,\n",
    "    data_collator=data_collator\n",
    ")\n",
    "\n",
    "# 开始训练\n",
    "trainer.train()\n",
    "trainer.save_model(\"./zhouyu_think_model/model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5c758140081e417c",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-07-23T01:26:47.307728Z",
     "start_time": "2025-07-23T01:26:45.445464Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "<|zhouyu_start|>user\n",
      "你叫什么名字<|zhouyu_end|>\n",
      "<|zhouyu_start|>assistant\n",
      "<think>\n",
      "我得思考一下\n",
      "</think>\n",
      "我叫周瑜\n",
      "<|zhouyu_end|>\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "chat_template = \"\"\"\n",
    "<|zhouyu_start|>user\n",
    "{question}<|zhouyu_end|>\n",
    "<|zhouyu_start|>assistant\n",
    "<think>\"\"\"\n",
    "\n",
    "prompt = \"你叫什么名字\"\n",
    "\n",
    "text = chat_template.format(question=prompt)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps\")\n",
    "model_inputs = tokenizer([text], return_tensors=\"pt\")\n",
    "model_inputs = model_inputs.to(device)\n",
    "generated_ids = model.generate(**model_inputs, max_new_tokens=200, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)\n",
    "content = tokenizer.decode(generated_ids[0])\n",
    "print(content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
